# Title     : TODO
# Objective : TODO
# Created by: Administrator
# Created on: 2019/6/18

createWhenNoExist <- function(f) {
  !dir.exists(f) && dir.create(f)
}

library(gplots)
library(Hmisc)
library(optparse)
library(ropls)
library(magrittr)
library(ComplexHeatmap)
library(tools)
library(circlize)
library(tidyverse)
library(extrafont)

option_list <- list(
  make_option("--i", default = "AllMet.csv", type = "character", help = "metabolite data file"),
  make_option("--g", default = "SampleInfo.csv", type = "character", help = "sample group file"),
  make_option("--et", default = "", type = "character", help = "extra data dir"),
  make_option("--rf", default = "r_Matrix.csv", type = "character", help = "coe value file"),
  make_option("--pf", default = "P_Matrix.csv", type = "character", help = "p value file"),
  make_option("--o", default = "Correlation_Heatmap.pdf", type = "character", help = "output pdf file"),
  make_option("--pc", default = "", type = "character", help = "plot config file")

)
opt <- parse_args(OptionParser(option_list = option_list))
font_import(paths = c("/usr/share/fonts/myFonts"), recursive = F, prompt = F)

args <- commandArgs(trailingOnly = F)
scriptPath <- dirname(sub("--file=", "", args[grep("--file", args)]))
source(str_c(scriptPath, "/parse_config_json.R"))
fontFamily <- plotConfigJson$font$fontFamily
baseFamily <- fontFamily

parent <- "./"
createWhenNoExist(parent)

finalParent <- paste0("./")
if (!dir.exists(finalParent)) {
  next
}

corFile <- str_c(finalParent, "/", opt$rf)

corData <- read_csv(corFile)

pData <- read_csv(str_c(finalParent, "/", opt$pf)) %>%
  column_to_rownames("X1") %>%
  as.matrix()

fileName <- str_c(finalParent, "/", opt$o)

plotData <- corData %>%
  column_to_rownames("X1") %>%
  as.matrix()

rowNum <- nrow(plotData)
colNum <- ncol(plotData)

row_names = rownames(plotData)
col_names = colnames(plotData)
max_row_chr_num = max(nchar(row_names))
max_col_chr_num = max(nchar(col_names))

defaultHeight <- max(2 + (rowNum - 10) * 0.1, 2) + 1 / 4 * max_col_chr_num
defaultWidth <- max(2 + (colNum - 10) * 0.1, 2) + 1 / 4 * max_row_chr_num
maxLength <- max(c(defaultWidth, defaultHeight))

defaultWidth <- maxLength
defaultHeight <- maxLength

plotConfigJson <- fromJSON(opt$pc)

finalWidth <- getFinalWidth(width, defaultWidth, opt$pc)
finalHeight <- getFinalHeight(height, defaultHeight, opt$pc)
finalColors <- getfinalColors(plotConfigJson)
naColor <- plotConfigJson$color$naColor
showRowName <- plotConfigJson$font$showRowName %>%
  as.logical()
showColName <- plotConfigJson$font$showColName %>%
  as.logical()
yRotate <- plotConfigJson$font$yRotate %>%
  as.numeric()
xRotate <- plotConfigJson$font$xRotate %>%
  as.numeric()
rowTreeHeight <- plotConfigJson$layout$rowTreeHeight %>%
  as.numeric()
colTreeHeight <- plotConfigJson$layout$colTreeHeight %>%
  as.numeric()
rowCluster <- plotConfigJson$layout$rowCluster %>%
  as.logical()
rowClusterMethod <- plotConfigJson$layout$rowClusterMethod
colCluster <- plotConfigJson$layout$colCluster %>%
  as.logical()
colClusterMethod <- plotConfigJson$layout$colClusterMethod

pdf(fileName, width = finalWidth, height = finalHeight)

eachHeight <- 1 / rowNum
colF <- colorRamp2(seq(-1, 1, length.out = length(finalColors)), finalColors)
htList <- Heatmap(plotData, col = colF, show_column_names = showColName, cluster_rows = rowCluster,
                  cluster_columns = colCluster, clustering_method_columns = colClusterMethod,
                  clustering_method_rows = rowClusterMethod, na_col = naColor, row_names_rot = yRotate,
                  row_names_gp = gpar(fontsize = yFont, fontfamily = baseFamily), show_row_names = showRowName,
                  show_heatmap_legend = F, column_names_rot = xRotate,
                  column_names_gp = gpar(fontsize = xFont, fontfamily = baseFamily),
                  row_dend_width = unit(rowTreeHeight, "mm"), column_dend_height = unit(colTreeHeight, "mm"),
                  cell_fun = function(j, i, x, y, width, height, fill) {
                    value <- pData[i, j]
                    str <- if (value < 0.01) {
                      "+"
                    }else if (value < 0.05) {
                      "*"
                    }else ""
                    top <- if (str == "*") {
                      y - unit(eachHeight * 0.12, "npc")
                    }else y
                    grid.text(str, x, top, gp = gpar(fontsize = 9), hjust = 0.5, vjust = 0.5, just = "center")
                  })

finalLegendTitle <- if (!is.na(legendTitle)) {
  legendTitle
}else {
  ""
}

rLgd <- Legend(col_fun = colF, at = seq(-1, 1, 0.5),
               title = finalLegendTitle,
               title_gp = gpar(fontsize = legendTitleFont, fontfamily = baseFamily), grid_height = unit(7.5, "mm"),
               labels_gp = gpar(fontsize = legendFont, fontfamily = baseFamily))

pLegend <- Legend(pch = c("+", "*"), type = "points", labels = c("p<0.01", "0.01<=p<0.05"), title = "",
                  labels_gp = gpar(fontsize = 9, fontfamily = baseFamily),
                  title_gp = gpar(fontsize = 12, fontfamily = baseFamily),
                  ncol = 1, grid_height = unit(5, "mm"),
)
pd <- packLegend(list = list(rLgd, pLegend), row_gap = unit(0.5, "cm"))

draw(htList, heatmap_legend_side = "right", annotation_legend_side = "right",
     padding = unit(c(0.5, 0.5, 0.5, 0.5), "cm"), annotation_legend_list = pd
)
dev.off()





